// RetriverSlave.cpp : implementation file
//

#include "stdafx.h"
#include "../RefreeRetrieval.h"
#include "RetriverSlave.h"
#include "cv.h"
#include "highgui.h"

#pragma comment(lib, "gdiplus")
#define GETARRAY(a, x, y, w)	((a)[(y)*(w) + (x)])
#define SETARRAY(val, a, x, y, w)	(a)[(y)*(w) + (x)] = (val)

extern bool SIFT_SVM_Predict(IplImage* img, char* kmeansModel, char* svmModel);

// CRetriverSlave

IMPLEMENT_DYNCREATE(CRetriverSlave, CWinThread)

BEGIN_MESSAGE_MAP(CRetriverSlave, CWinThread)
	ON_THREAD_MESSAGE(CRetrieverMaster::WM_SLAVEMESSAGE, OnWmSlaveMessage)
END_MESSAGE_MAP()

CRetriverSlave::CRetriverSlave()
	: m_Master(NULL),
	m_bContinue(TRUE),
	m_clusterCodeBook(NULL),
	m_iPatchSize(16),
	m_patchMode(PatchModeNonOverlapped),
	m_iRegionThreshold(4),
	m_fntFrameIdx(new Font(L"Tahoma", 9, FontStyleBold)),
	m_brFrameIdx(new SolidBrush(Color::Red))
{
	if(GetEncoderClsid(L"image/png", &m_pngClsid) == -1)
	{
		ASSERT(FALSE);
	}

	const int BUFSIZE = 1024;
	DWORD dwBufSize=BUFSIZE, dwRetVal;
	UINT uRetVal;
	TCHAR szTempFileName[BUFSIZE];  
	TCHAR lpTempPathBuffer[BUFSIZE];

	dwRetVal = GetTempPath(dwBufSize, lpTempPathBuffer);
	ASSERT(dwRetVal <= dwBufSize && (dwRetVal != 0));
	
	uRetVal = GetTempFileName(lpTempPathBuffer, TEXT("REFREE"),
		0, szTempFileName);
	ASSERT(uRetVal != 0);

	m_sTempFile = CString(szTempFileName);
	CString2Char(m_sTempFile, m_sAnsiTempFile);
}

CRetriverSlave::~CRetriverSlave()
{
	delete m_fntFrameIdx;
	delete m_brFrameIdx;
}

/*=============================================================*/

/*=============================================================*/

BOOL CRetriverSlave::InitInstance()
{
	return TRUE;
}

int CRetriverSlave::ExitInstance()
{
	return CWinThread::ExitInstance();
}

void CRetriverSlave::SetMaster(CRetrieverMaster* master)
{
	m_Master = master;
}

void CRetriverSlave::SetParams(CCodeBook* codeBook, CFeatureExtract* extractor, 
				int iPatchSize, PatchMode patchMode, int regionThreshold,
				char* siftKmeansModel, char* siftSVMModel)
{
	m_clusterCodeBook = codeBook;
	m_extractor = extractor;
	m_iPatchSize = iPatchSize;
	m_patchMode = patchMode;
	m_iRegionThreshold = regionThreshold;
	m_sSIFTKmeanModel = siftKmeansModel;
	m_sSIFTSVMModel = siftSVMModel;
}

/*=============================================================*/

/*=============================================================*/

BOOL CRetriverSlave::FindCandicate(CArray<CArray<CPoint>*>& arr, int x, int y, CArray<int>& grps)
{
	grps.RemoveAll();
	for(INT_PTR i = arr.GetCount() - 1; i >= 0; --i)
	{
		CArray<CPoint>* a = arr[i];
		for(INT_PTR j = a->GetCount() - 1; j >= 0; --j)
		{
			CPoint p = a->ElementAt(j);
			if(p.x == x && p.y == y)
			{
				grps.Add(i);
			}
		}
	}
	return (grps.GetCount() > 0);
}

Rect CRetriverSlave::GetBoundingBox(CArray<CPoint>* candicates)
{
	int iMinX, iMaxX, iMinY, iMaxY;
	iMaxX = iMaxY = -1;
	iMinX = iMinY = 999999;

	for(INT_PTR i = candicates->GetCount() - 1; i >=0 ; --i)
	{
		CPoint pt = candicates->GetAt(i);
		pt.x *= m_iPatchSize;
		pt.y *= m_iPatchSize;

		if(pt.x < iMinX)	iMinX = pt.x;
		if(pt.y < iMinY)	iMinY = pt.y;
		if((pt.x + m_iPatchSize) > iMaxX)	iMaxX = pt.x + m_iPatchSize;
		if((pt.y + m_iPatchSize) > iMaxY)	iMaxY = pt.y + m_iPatchSize;
	}
	return Rect(iMinX, iMinY, iMaxX - iMinX, iMaxY - iMinY);
}

int CRetriverSlave::GetEncoderClsid(const WCHAR* format, CLSID* pClsid)
{
	UINT  num = 0;          // number of image encoders
	UINT  size = 0;         // size of the image encoder array in bytes

	ImageCodecInfo* pImageCodecInfo = NULL;

	GetImageEncodersSize(&num, &size);
	if(size == 0)
		return -1;  // Failure

	pImageCodecInfo = (ImageCodecInfo*)(malloc(size));
	if(pImageCodecInfo == NULL)
		return -1;  // Failure

	GetImageEncoders(num, size, pImageCodecInfo);

	for(UINT j = 0; j < num; ++j)
	{
		if( wcscmp(pImageCodecInfo[j].MimeType, format) == 0 )
		{
			*pClsid = pImageCodecInfo[j].Clsid;
			free(pImageCodecInfo);
			return j;  // Success
		}    
	}

	free(pImageCodecInfo);
	return -1;  // Failure
}

void CRetriverSlave::CString2Char(CString& s, char arr[])
{
	size_t tmp;
	wcstombs_s(&tmp, arr, MAX_PATH * sizeof(char),
		s.GetBuffer(), MAX_PATH*sizeof(char));
	s.ReleaseBuffer();
}

BOOL CRetriverSlave::CheckLowFeature(Bitmap& bmp)
{
	if(bmp.Save(m_sTempFile, &m_pngClsid) != Ok)
		ASSERT(FALSE);

	IplImage* img = cvLoadImage(m_sAnsiTempFile);
	bool b = SIFT_SVM_Predict(img, m_sSIFTKmeanModel, m_sSIFTSVMModel);
	cvReleaseImage(&img);

	return (b == true);
}

void CRetriverSlave::ProcessFrame(BitmapFrame* frame)
{
	const int MAX_REGIONCOUNT = 50;
	Bitmap bmpSrc(frame->Frame, NULL);
	Graphics grpSrc(&bmpSrc);

	// draw frame index
	CString sFrameCount;
	StringFormat strFmt;
	strFmt.SetAlignment(StringAlignmentNear);
	strFmt.SetFormatFlags(StringFormatFlagsNoClip);
	strFmt.SetLineAlignment(StringAlignmentCenter);

	sFrameCount.Format(L"%d", frame->FrameIndex);
	grpSrc.DrawString(sFrameCount, sFrameCount.GetLength(), m_fntFrameIdx, 
		RectF(0, 0, 100, 30), NULL, m_brFrameIdx);

	Bitmap bmpPatch(m_iPatchSize, m_iPatchSize, PixelFormat32bppARGB);
	Graphics grpPatch(&bmpPatch);
	grpPatch.SetSmoothingMode(SmoothingModeAntiAlias);
	
	Rect rcDest(0, 0, m_iPatchSize, m_iPatchSize);
	CArray<double> vecFeature;
	CArray<CArray<CPoint>*> arrCandicates;
	CArray<int> grp;
	byte* arrMark;
	int iStepWidth, iStepHeight;
	int iPatchCntWidth, iPatchCntHeight;

	switch(m_patchMode)
	{
	case PatchModeHalfOverlapped:
		iStepWidth = iStepHeight = m_iPatchSize / 2;
		break;
	case PatchModeNonOverlapped:
		iStepWidth = iStepHeight = m_iPatchSize;
		break;
	case PatchModePerPixel:
		iStepWidth = iStepHeight = 1;
		break;
	}
	iPatchCntWidth = (bmpSrc.GetWidth() - m_iPatchSize + 1) / iStepWidth;
	iPatchCntHeight = (bmpSrc.GetHeight() - m_iPatchSize + 1) / iStepHeight;

	arrMark = new byte[iPatchCntWidth * iPatchCntHeight];
	memset(arrMark, 0, iPatchCntWidth * iPatchCntHeight * sizeof(byte));

	// find candicates
	for(int i = 0; i < iPatchCntWidth; ++i)
	{
		for(int j = 0; j < iPatchCntHeight; ++j)
		{
			// paint to the patch
			grpPatch.DrawImage(&bmpSrc, rcDest, i * iStepWidth, j * iStepHeight, 
				m_iPatchSize, m_iPatchSize, UnitPixel);

			// calculate feature vector
			m_extractor->Calc(&bmpPatch, vecFeature);

			// find it cluster
			INT iCluster = m_clusterCodeBook->FindCluster(vecFeature);

			if(iCluster != -1)
			{
				SETARRAY(1, arrMark, i, j, iPatchCntWidth);
				
				BOOL bAdded = FALSE;

				// check left
				if(i > 0 && GETARRAY(arrMark, i - 1, j, iPatchCntWidth) > 0)
				{
					if(FindCandicate(arrCandicates, i-1, j, grp))
					{
						for(INT_PTR k = grp.GetCount() - 1; k >= 0; --k)
						{
							if(arrCandicates[grp[k]]->GetCount() < MAX_REGIONCOUNT)
							{
								arrCandicates[grp[k]]->Add(CPoint(i, j));
								bAdded = TRUE;
							}
						}
					}
				}

				// check top
				if(j > 0 && GETARRAY(arrMark, i, j - 1, iPatchCntWidth) > 0)
				{
					if(FindCandicate(arrCandicates, i, j - 1, grp))
					{
						for(INT_PTR k = grp.GetCount() - 1; k >= 0; --k)
						{
							if(arrCandicates[grp[k]]->GetCount() < MAX_REGIONCOUNT)
							{
								arrCandicates[grp[k]]->Add(CPoint(i, j));
								bAdded = TRUE;
							}
						}
					}				
				}

				if(!bAdded)
				{
					CArray<CPoint>* arrNewGroup = new CArray<CPoint>();
					arrNewGroup->Add(CPoint(i, j));
					arrCandicates.Add(arrNewGroup);
				}
			}
		}
	}

	for(INT_PTR i = arrCandicates.GetCount() - 1; i >= 0; --i)
	{
		CArray<CPoint>* arrGrp = arrCandicates[i];
		if(arrGrp->GetCount() > m_iRegionThreshold)
		{
			Rect rcHot = GetBoundingBox(arrGrp);
			Bitmap bmpHot(rcHot.Width, rcHot.Height, PixelFormat32bppARGB);
			Graphics grpHot(&bmpHot);
			grpHot.DrawImage(&bmpSrc, RectF(0, 0, (REAL)rcHot.Width, (REAL)rcHot.Height),
				(REAL)rcHot.X, (REAL)rcHot.Y, (REAL)rcHot.Width, (REAL)rcHot.Height,
				UnitPixel);
			if(CheckLowFeature(bmpHot))
			{
				// draw a read rectangle on bitmap
				grpSrc.DrawRectangle(&Pen(Color::Red, 1.5f), rcHot);
			}
		}

		delete arrGrp;
	}
	arrCandicates.RemoveAll();
	delete[] arrMark;

	::DeleteObject(frame->Frame);
	bmpSrc.GetHBITMAP(Color(0, 0, 0), &(frame->Frame));
}

int CRetriverSlave::Run()
{
	MSG msg;
	
	m_Master->PostThreadMessage(CRetrieverMaster::WM_MASTERMESSAGE, 
		CRetrieverMaster::MasterWParamJobRequested, (LPARAM)this);

	while(m_bContinue)
	{
	
		while(PeekMessage(&msg, (HWND)-1, 0, 0, PM_REMOVE))
		{
			switch(msg.message)
			{
			case CRetrieverMaster::WM_SLAVEMESSAGE:
				OnWmSlaveMessage(msg.wParam, msg.lParam);
				break;
			default:
				TRACE(L"Slave #%d: Error: cannot dispatch message #%d\r\n", m_nThreadID, msg.message);
			}
		}
	}

	TRACE(L"Slave #%d: Fininshed.\r\n", m_nThreadID);
	return CWinThread::Run();
}

/*================================================================*/
// CRetriverSlave message handlers
/*================================================================*/

void CRetriverSlave::OnWmSlaveMessage(WPARAM wParam, LPARAM lParam)
{
	switch((UINT)wParam)
	{
	case CRetrieverMaster::SlaveWParamNewBitmapArrived:
		{
			BitmapFrame* frame = (BitmapFrame*)lParam;

			TRACE(L"Slave #%d: Received new bitmap: %d\r\n", m_nThreadID, frame->FrameIndex);

			ProcessFrame(frame);

			TRACE(L"Slave #%d: End processing frame #%d.\r\n", m_nThreadID, frame->FrameIndex);

			m_Master->PostThreadMessage(CRetrieverMaster::WM_MASTERMESSAGE,
				CRetrieverMaster::MasterWParamJobDone,
				(LPARAM)frame);

			TRACE(L"Slave #%d: Sending request to master.\r\n", m_nThreadID);

			m_Master->PostThreadMessage(CRetrieverMaster::WM_MASTERMESSAGE,
				CRetrieverMaster::MasterWParamJobRequested,
				(LPARAM)this);
		}
		break;
	case CRetrieverMaster::SlaveWParamJobEnded:
		{
			m_bContinue = FALSE;
			TRACE(L"Slave #%d: Received job END.\r\n", m_nThreadID);

			m_Master->PostThreadMessage(CRetrieverMaster::WM_MASTERMESSAGE,
				CRetrieverMaster::MasterWParamSlaveEnded, 0);
		}
		break;
	}
}

#undef GETARRAY
#undef SETARRAY
